import { Key } from "react-aria-components";
import { Controller, useFormContext } from "react-hook-form";

import { Annotation, AnnotationConfig } from "@phoenix/components/annotation";
import { CategoricalAnnotationInput } from "@phoenix/components/annotation/CategoricalAnnotationInput";
import { ContinuousAnnotationInput } from "@phoenix/components/annotation/ContinuousAnnotationInput";
import { FreeformAnnotationInput } from "@phoenix/components/annotation/FreeformAnnotationInput";

export type AnnotationFormData = {
  name: string;
  score?: number | null;
  label?: string | null;
  explanation?: string | null;
};

export type SpanAnnotationInputProps = {
  /**
   * The annotation config to that represents the type of annotation input to render
   */
  annotationConfig: AnnotationConfig;

  /**
   * The annotation to populate the form with
   */
  annotation?: Annotation;
};

/**
 * A form to create or edit a span annotation
 */
export function SpanAnnotationInput(props: SpanAnnotationInputProps) {
  const { annotationConfig, annotation } = props;
  const { control } = useFormContext<Record<string, Annotation>>();

  return (
    <div>
      {annotationConfig.annotationType === "CATEGORICAL" && (
        <Controller
          control={control}
          name={annotationConfig.name}
          render={({ field: { value: _value, ...field } }) => (
            <CategoricalAnnotationInput
              annotationConfig={annotationConfig}
              annotation={annotation}
              {...field}
              selectedKey={_value?.label}
              onSubmitExplanation={(explanation) => {
                if (annotation?.id) {
                  field.onChange({
                    ...annotation,
                    name: annotationConfig.name,
                    explanation,
                  });
                }
              }}
              onSelectionChange={(_selectedKey) => {
                let selectedKey: Key | null = _selectedKey;
                if (selectedKey === _value?.label) {
                  selectedKey = null;
                }
                if (typeof selectedKey === "string" && selectedKey != null) {
                  const newAnnotation: Annotation = {
                    ...annotation,
                    id: annotation?.id,
                    name: annotationConfig.name,
                    label: selectedKey,
                    score:
                      annotationConfig.values?.find(
                        (value) => value.label === selectedKey
                      )?.score ?? null,
                  };
                  field.onChange(newAnnotation);
                } else {
                  field.onChange({
                    ...annotation,
                    id: annotation?.id,
                    name: annotationConfig.name,
                    label: null,
                    score: null,
                  });
                }
              }}
            />
          )}
        />
      )}
      {annotationConfig.annotationType === "CONTINUOUS" && (
        <Controller
          control={control}
          name={annotationConfig.name}
          render={({ field: { value: _value, ...field } }) => (
            <ContinuousAnnotationInput
              annotationConfig={annotationConfig}
              annotation={annotation}
              {...field}
              onSubmitExplanation={(explanation) => {
                if (annotation?.id) {
                  field.onChange({
                    ...annotation,
                    name: annotationConfig.name,
                    explanation,
                  });
                }
              }}
              value={_value?.score ?? undefined}
              onChange={(value) => {
                field.onChange({
                  ...annotation,
                  id: annotation?.id,
                  name: annotationConfig.name,
                  score: value,
                });
              }}
            />
          )}
        />
      )}
      {annotationConfig.annotationType === "FREEFORM" && (
        <Controller
          control={control}
          name={annotationConfig.name}
          render={({ field: { value: _value, ...field } }) => (
            <FreeformAnnotationInput
              annotationConfig={annotationConfig}
              annotation={annotation}
              {...field}
              value={_value?.explanation ?? ""}
              onChange={(value) => {
                field.onChange({
                  ...annotation,
                  id: annotation?.id,
                  name: annotationConfig.name,
                  explanation: value,
                });
              }}
            />
          )}
        />
      )}
    </div>
  );
}
